import numpy as np
import matplotlib.pyplot as plt
from PIL import Image

# 读取彩色图像并转换为 float32
img = Image.open('D:\\2D-DFT\image\\testhaze2.png').convert('RGB')
img_np = np.array(img, dtype=np.float32)

# 创建结果图像容器
perturbed_img = np.zeros_like(img_np)

# 提取图像大小
H, W = img_np.shape[:2]

# 预生成统一的扰动因子（对所有通道统一）
# 比如扰动幅度为 ±20%
random_scale = np.random.uniform(0.4, 1.6, size=(H, W))

# 对 RGB 三个通道分别处理
for c in range(3):
    channel = img_np[:, :, c]

    # 傅里叶变换
    f = np.fft.fft2(channel)
    mag = np.abs(f)
    phase = np.angle(f)

    # 使用统一扰动因子扰动幅度谱
    perturbed_mag = mag * random_scale

    # 重建频谱 + 逆变换
    f_perturbed = perturbed_mag * np.exp(1j * phase)
    new_channel = np.fft.ifft2(f_perturbed)
    new_channel = np.real(new_channel)
    new_channel = np.clip(new_channel, 0, 255)

    # 放回结果图像
    perturbed_img[:, :, c] = new_channel

# 转成 uint8 显示
perturbed_img = perturbed_img.astype(np.uint8)

# 显示原图与扰动图
plt.figure(figsize=(10, 4))
plt.subplot(1, 2, 1)
plt.title("Original RGB Image")
plt.imshow(img_np.astype(np.uint8))
plt.axis('off')

plt.subplot(1, 2, 2)
plt.title("Unified Magnitude Perturbation (RGB)")
plt.imshow(perturbed_img)
plt.axis('off')

plt.tight_layout()
plt.show()
